import datetime
from concurrent.futures import ThreadPoolExecutor
from typing import Dict, List, Optional

import pandas as pd
import pytz
from google.protobuf.json_format import MessageToDict
from tornado.concurrent import run_on_executor

from coin.proto.coin_request_pb2 import (
    AccountInfoProto,
    StrategyRequestProto,)
from coin.proto.coin_strategy_summary_pb2 import StrategyStatLog
from coin.tool.strat_monitor.plot.strat_plotter import (DailyStrategyPlotterConfig,
                                                        StrategyPlotInfo,
                                                        StrategyPlotMsg,
                                                        SymbolPlotMsg,)
from coin.support.proto_log.logic.log_info import StratInfo
from coin.support.pta.logging.strategy_stat_logger import (
    gen_strategy_stat_dir,
    StrategyStatReader,)


class _Onlog:
  def __init__(self):
    self.strat_stat_logs = []

  def on_log(self, timestamp: int, log):
    pb = StrategyStatLog()
    pb.ParseFromString(log)
    self.strat_stat_logs.append(pb)


def convert_to_strategy_plot_info(history):
  strategy_plot_infos = []
  for entry in history:
    strat_history = entry['strat_history']
    symbol_history = entry['symbol_history']
    for strat_elem in strat_history:
      accounting_currency = strat_elem['accounting_currency']
      strat_plot_msg = StrategyPlotMsg(strategy_name=strat_elem['strategy_name'],
                                       machine=strat_elem['machine'],
                                       accounting_currency=accounting_currency,
                                       pnl=strat_elem['pnl'])
      symbol_plot_msgs = []
      for symbol_elem in symbol_history:
        if accounting_currency == symbol_elem['accounting_currency']:
          symbol_plot_msg = SymbolPlotMsg(account_id=symbol_elem['account_id'],
                                          symbol=symbol_elem['symbol'],
                                          pnl=symbol_elem['pnl'],
                                          position=symbol_elem['position'])
          symbol_plot_msgs.append(symbol_plot_msg)
      strategy_plot_info = StrategyPlotInfo(strategy_plot_msg=strat_plot_msg,
                                            symbol_plot_msg_list=symbol_plot_msgs)
      strategy_plot_infos.append(strategy_plot_info)
  return strategy_plot_infos


class AsyncStrategyPlotExecutor(object):
  def __init__(self,
               trading_date: datetime.date,
               strat_stat_log_dir: str,
               plot_dir: str,
               account_infos: List[AccountInfoProto],
               pnl_adj: Optional[Dict[str, List[Dict]]]=None):
    self._trading_date = trading_date
    self._plot_dir = plot_dir
    self._strat_stat_log_dir = strat_stat_log_dir
    self._account_infos = account_infos
    self._acct_df = pd.DataFrame(
        [MessageToDict(acct_info, preserving_proto_field_name=True)
        for acct_info in account_infos]
    )
    self._pnl_adj = pnl_adj
    self._executor = ThreadPoolExecutor(max_workers=1)

  @property
  def trading_date(self):
    return self._trading_date

  @run_on_executor(executor='_executor')
  def do_pnl_plot(self,
                  strat_infos: List[StratInfo]):
    from coin.support.pta.plot.strat_plotter import StrategyPlotter
    plotter = StrategyPlotter(self._acct_df, self._pnl_adj)
    start_dt = datetime.datetime.combine(
        self._trading_date, datetime.time.min.replace(tzinfo=pytz.UTC))
    end_dt = datetime.datetime.combine(
        self._trading_date, datetime.time.max.replace(tzinfo=pytz.UTC))
    for strat_info in strat_infos:
      onlog = _Onlog()
      stat_reader = StrategyStatReader()
      stat_reader.run_from_archive(
          on_log_callback=onlog.on_log,
          start_time=start_dt,
          end_time=end_dt,
          root_dir=self._strat_stat_log_dir,
          machine=strat_info.machine,
          strategy_name=strat_info.strategy_name,
      )
      if len(onlog.strat_stat_logs) == 0:
        continue
      plot_dir = gen_strategy_stat_dir(
          log_root=self._plot_dir,
          trading_date=start_dt,
          strat_request=StrategyRequestProto(strategy_name=strat_info.strategy_name),
          machine=strat_info.machine)
      plotter.do_pnl_plot(
          strat_name=strat_info.strategy_name,
          strat_stat_logs=onlog.strat_stat_logs,
          start_time=start_dt,
          end_time=end_dt,
          plot_dir=plot_dir)

  @run_on_executor(executor='_executor')
  def do_order_plot(self, plot_info):
    from coin.tool.strat_monitor.plot.strat_plotter import StrategyPlotter
    config = DailyStrategyPlotterConfig(
        self._trading_date, self._trading_date, self._plot_dir)
    plotter = StrategyPlotter(config, self._account_infos)
    plotter.do_order_plot(plot_info)
